import json
import random
import time
import re
import os
from tqdm import tqdm
import logging
from langchain_openai import ChatOpenAI
from typing import List, Dict, Tuple, Optional

# 设置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


class SelfInstructGenerator:
    def __init__(self,
                 api_key: str,
                 base_url: Optional[str] = None,
                 model_name: str = "gpt-3.5-turbo",
                 seed_tasks_path: Optional[str] = None):
        """
        初始化 Self-Instruct 生成器

        Args:
            api_key: OpenAI API key
            base_url: 自定义 API 端点（用于本地部署或第三方代理）
            model_name: 使用的模型名称
            seed_tasks_path: 种子任务文件路径
        """
        self.client = ChatOpenAI(
        model="qwen-plus",  # 或者 "qwen-turbo", "qwen-max" 等
        openai_api_key='sk-965dc39b016c49ecbe29de180f4db2b6',
        openai_api_base="https://dashscope.aliyuncs.com/compatible-mode/v1",  # 兼容模式端点
        temperature=0.7,
        max_tokens=2048,
    )
        self.model_name = model_name
        self.seed_tasks = self._load_seed_tasks(seed_tasks_path)

        # 指令生成提示模板
        self.instruction_generation_prompt = """请生成5个多样化的任务指令。这些指令将被提供给GPT模型，我们将评估GPT模型执行指令的能力。

要求：
1. 尽量使指令多样化（包括生成类型、改写类型、分类类型等）
2. 指令应该是自包含的，不需要额外的上下文
3. 指令应该清晰明确
4. 避免重复或相似的指令

生成的指令："""

        # 分类任务识别提示
        self.classification_prompt = """请判断以下指令是否是分类任务。分类任务是指将输入分为两个或多个类别的任务。

指令：{instruction}

请只回答"是"或"否"："""

        # 输入-输出生成提示
        self.instance_generation_prompt = """请为以下指令生成2个对应的输入和输出对。

指令：{instruction}

要求：
1. 输入应该多样化且有意义
2. 输出应该正确执行指令
3. 对于分类任务，请确保输出是类别标签

请以JSON格式输出：
{{
    "instances": [
        {{
            "input": "输入1",
            "output": "输出1"
        }},
        {{
            "input": "输入2", 
            "output": "输出2"
        }}
    ]
}}"""

    def _load_seed_tasks(self, seed_tasks_path: Optional[str]) -> List[Dict]:
        """加载种子任务"""
        if seed_tasks_path and os.path.exists(seed_tasks_path):
            with open(seed_tasks_path, 'r', encoding='utf-8') as f:
                return json.load(f)
        else:
            # 默认种子任务
            return [
                {
                    "instruction": "写一首关于春天的诗",
                    "is_classification": False,
                    "instances": [
                        {"input": "", "output": "春风轻拂面，花开满园香。燕子归巢来，生机盎然长。"}
                    ]
                },
                {
                    "instruction": "将以下英文翻译成中文",
                    "is_classification": False,
                    "instances": [
                        {"input": "Hello, how are you?", "output": "你好，你怎么样？"}
                    ]
                },
                {
                    "instruction": "判断电影评论的情感倾向",
                    "is_classification": True,
                    "instances": [
                        {"input": "这部电影太精彩了，演员表演出色！", "output": "正面"}
                    ]
                },
                {
                    "instruction": "总结以下文章的主要内容",
                    "is_classification": False,
                    "instances": [
                        {"input": "人工智能正在改变世界...", "output": "文章讨论了人工智能对世界的影响。"}
                    ]
                }
            ]

    def call_openai_api(self, prompt: str, max_tokens: int = 1000) -> str:
        """调用 OpenAI API（兼容新版本）"""
        try:
            response = self.client.invoke(prompt)
            return response.content.strip()
        except Exception as e:
            logger.error(f"API调用失败: {e}")
            time.sleep(10)  # 遇到限制时等待
            return ""

    def generate_instructions(self, num_instructions: int = 50) -> List[str]:
        """生成指令"""
        logger.info("开始生成指令...")
        instructions = set()

        # 初始使用种子指令
        for task in self.seed_tasks:
            instructions.add(task["instruction"])

        # 生成新指令
        while len(instructions) < num_instructions:
            response = self.call_openai_api(self.instruction_generation_prompt)

            if response:
                # 解析生成的指令
                new_instructions = self._parse_instructions(response)
                for instr in new_instructions:
                    if instr and len(instr) > 10 and instr not in instructions:  # 简单过滤
                        instructions.add(instr)
                        logger.info(f"新指令: {instr}")

            logger.info(f"已生成 {len(instructions)}/{num_instructions} 条指令")
            time.sleep(1)  # 避免API限制

            # 如果生成困难，提前退出
            if len(instructions) >= num_instructions * 0.8 and len(instructions) < num_instructions:
                # 复制一些现有指令并修改
                existing_instructions = list(instructions)
                for existing_instr in random.sample(existing_instructions, min(5, len(existing_instructions))):
                    modified_instr = self._modify_instruction(existing_instr)
                    if modified_instr and modified_instr not in instructions:
                        instructions.add(modified_instr)

        return list(instructions)[:num_instructions]

    def _parse_instructions(self, response: str) -> List[str]:
        """解析生成的指令"""
        instructions = []
        lines = response.split('\n')

        for line in lines:
            line = line.strip()
            # 移除编号
            line = re.sub(r'^\d+[\.\)]\s*', '', line)
            # 移除引号
            line = re.sub(r'^["\']|["\']$', '', line)
            if line and len(line) > 10 and not line.startswith('请'):
                instructions.append(line)

        return instructions

    def _modify_instruction(self, instruction: str) -> str:
        """修改现有指令以生成变体"""
        prompt = f"""请对以下指令进行修改，创建一个相似但不同的新指令：

原指令：{instruction}

新指令："""

        response = self.call_openai_api(prompt, max_tokens=100)
        return response.strip() if response else ""

    def is_classification_task(self, instruction: str) -> bool:
        """判断是否为分类任务"""
        prompt = self.classification_prompt.format(instruction=instruction)
        response = self.call_openai_api(prompt, max_tokens=10)

        return "是" in response or "yes" in response.lower()

    def generate_instances(self, instruction: str, is_classification: bool) -> List[Dict]:
        """为指令生成输入-输出实例"""
        prompt = self.instance_generation_prompt.format(instruction=instruction)
        response = self.call_openai_api(prompt)

        if not response:
            return []

        instances = self._parse_instances_from_response(response, instruction, is_classification)

        # 如果解析失败，使用备选方法
        if not instances:
            instances = self._generate_instances_fallback(instruction, is_classification)

        return instances

    def _parse_instances_from_response(self, response: str, instruction: str, is_classification: bool) -> List[Dict]:
        """从响应中解析实例"""
        instances = []

        # 方法1: 尝试解析JSON
        try:
            # 查找JSON结构
            json_match = re.search(r'\{.*\}', response, re.DOTALL)
            if json_match:
                json_str = json_match.group()
                data = json.loads(json_str)
                instances_data = data.get("instances", [])

                for instance_data in instances_data:
                    instances.append({
                        "instruction": instruction,
                        "input": instance_data.get("input", ""),
                        "output": instance_data.get("output", ""),
                        "is_classification": is_classification
                    })
                return instances
        except json.JSONDecodeError:
            pass

        # 方法2: 手动解析键值对
        lines = response.split('\n')
        current_instance = {}

        for line in lines:
            line = line.strip()
            if 'input' in line.lower() and ':' in line:
                current_instance['input'] = line.split(':', 1)[1].strip().strip('"\'')
            elif 'output' in line.lower() and ':' in line:
                current_instance['output'] = line.split(':', 1)[1].strip().strip('"\'')

                if current_instance.get('input') is not None and current_instance.get('output'):
                    instances.append({
                        "instruction": instruction,
                        "input": current_instance['input'],
                        "output": current_instance['output'],
                        "is_classification": is_classification
                    })
                    current_instance = {}

        return instances

    def _generate_instances_fallback(self, instruction: str, is_classification: bool) -> List[Dict]:
        """备用的实例生成方法"""
        instances = []

        # 生成单个实例的提示
        prompt = f"""请为以下指令生成一个输入和输出对：

指令：{instruction}

输入：
输出："""

        for _ in range(2):  # 生成2个实例
            response = self.call_openai_api(prompt)
            if response:
                lines = response.split('\n')
                input_line = None
                output_line = None

                for line in lines:
                    if line.startswith('输入：'):
                        input_line = line[3:].strip()
                    elif line.startswith('输出：'):
                        output_line = line[3:].strip()

                if input_line is not None and output_line:
                    instances.append({
                        "instruction": instruction,
                        "input": input_line,
                        "output": output_line,
                        "is_classification": is_classification
                    })

            time.sleep(0.5)

        return instances

    def filter_instances(self, instances: List[Dict]) -> List[Dict]:
        """过滤低质量实例"""
        filtered_instances = []

        for instance in instances:
            if self._is_high_quality(instance):
                filtered_instances.append(instance)

        logger.info(f"过滤后保留 {len(filtered_instances)}/{len(instances)} 个实例")
        return filtered_instances

    def _is_high_quality(self, instance: Dict) -> bool:
        """判断实例质量"""
        input_text = instance.get("input", "")
        output_text = instance.get("output", "")

        # 过滤条件
        if len(output_text) < 3:  # 输出太短
            return False
        if len(output_text) > 1000:  # 输出太长
            return False
        if input_text and len(input_text) < 2:  # 输入太短
            return False
        if "抱歉" in output_text or "无法" in output_text or "sorry" in output_text.lower():  # 模型拒绝回答
            return False
        if output_text.strip() == input_text.strip():  # 输入输出相同
            return False

        return True

    def remove_duplicate_instances(self, instances: List[Dict]) -> List[Dict]:
        """去除重复实例"""
        unique_instances = []
        seen = set()

        for instance in instances:
            # 基于指令和输入的哈希去重
            key = (instance["instruction"], instance.get("input", ""))
            if key not in seen:
                seen.add(key)
                unique_instances.append(instance)

        logger.info(f"去重后保留 {len(unique_instances)}/{len(instances)} 个实例")
        return unique_instances

    def generate_dataset(self,
                         num_instructions: int = 50,
                         output_path: str = "self_instruct_data.json") -> List[Dict]:
        """生成完整的数据集"""
        logger.info("开始生成 Self-Instruct 数据集...")

        all_instances = []

        # 1. 生成指令
        instructions = self.generate_instructions(num_instructions)
        logger.info(f"成功生成 {len(instructions)} 条指令")

        # 2. 为每条指令生成实例
        for i, instruction in enumerate(tqdm(instructions, desc="生成实例")):
            logger.info(f"处理指令 {i + 1}/{len(instructions)}: {instruction}")

            # 判断是否为分类任务
            is_classification = self.is_classification_task(instruction)
            logger.info(f"  分类任务: {is_classification}")

            # 生成实例
            instances = self.generate_instances(instruction, is_classification)
            logger.info(f"  生成 {len(instances)} 个实例")

            # 过滤实例
            filtered_instances = self.filter_instances(instances)

            all_instances.extend(filtered_instances)

            # 保存中间结果
            if (i + 1) % 10 == 0:
                temp_output = output_path.replace('.json', f'_temp_{i + 1}.json')
                with open(temp_output, 'w', encoding='utf-8') as f:
                    json.dump(all_instances, f, ensure_ascii=False, indent=2)
                logger.info(f"中间结果已保存到 {temp_output}")

            time.sleep(1)  # 避免API限制

        # 3. 去重
        all_instances = self.remove_duplicate_instances(all_instances)

        # 4. 保存最终数据集
        with open(output_path, 'w', encoding='utf-8') as f:
            json.dump(all_instances, f, ensure_ascii=False, indent=2)

        logger.info(f"数据集已保存到 {output_path}, 共 {len(all_instances)} 个实例")

        # 打印统计信息
        self._print_statistics(all_instances)

        return all_instances

    def _print_statistics(self, instances: List[Dict]):
        """打印数据集统计信息"""
        total_instances = len(instances)
        classification_tasks = sum(1 for inst in instances if inst.get("is_classification", False))
        instances_with_input = sum(1 for inst in instances if inst.get("input", "").strip())

        print("\n=== 数据集统计 ===")
        print(f"总实例数: {total_instances}")
        print(f"分类任务: {classification_tasks}")
        print(f"非分类任务: {total_instances - classification_tasks}")
        print(f"有输入的实例: {instances_with_input}")
        print(f"无输入的实例: {total_instances - instances_with_input}")

        # 平均长度
        if instances:
            avg_input_len = sum(len(inst.get("input", "")) for inst in instances) / total_instances
            avg_output_len = sum(len(inst.get("output", "")) for inst in instances) / total_instances
            print(f"平均输入长度: {avg_input_len:.1f} 字符")
            print(f"平均输出长度: {avg_output_len:.1f} 字符")


def main():
    # 初始化生成器
    generator = SelfInstructGenerator(
        api_key="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",  # 替换为您的API key
        model_name="gpt-3.5-turbo",
        # base_url="https://api.openai.com/v1"  # 如果需要自定义端点
    )

    # 生成小规模数据集进行测试
    dataset = generator.generate_dataset(
        num_instructions=10,  # 生成10条指令进行测试
        output_path="self_instruct_demo.json"
    )

    # 查看生成的数据样例
    print("\n=== 生成的数据样例 ===")
    for i, instance in enumerate(dataset[:3]):
        print(f"\n--- 样例 {i + 1} ---")
        print(f"指令: {instance['instruction']}")
        print(f"输入: {instance.get('input', 'N/A')}")
        print(f"输出: {instance['output']}")
        print(f"分类任务: {instance['is_classification']}")


if __name__ == "__main__":
    main()